'''
Author: cw
Date: 2025-08-19
Description: ...
'''
import logging
import sys
import asyncio
import traceback
import json
import uuid
from datetime import datetime

import websockets

from modules.task_mgr import ProcessTask
from modules import task_mgr


logging.basicConfig(level=logging.INFO, stream=sys.stdout)


class TaskClient:
    def __init__(self, server_url="ws://localhost:3001/ws"):
        self.server_url = server_url
        self.client_id = str(uuid.uuid4())
        self.websocket = None
        self.is_connected = False

    async def connect(self):
        """连接到服务器"""
        try:
            self.websocket = await websockets.connect(self.server_url)
            self.is_connected = True
            logging.info(f"客户端 {self.client_id} 已连接到服务器")
            return True
        except Exception as e:
            logging.error(f"连接失败: {e}")
            return False

    async def start_heartbeat(self):
        """启动心跳定时器"""
        while self.is_connected:
            await self.send_heartbeat()
            await asyncio.sleep(30)  # 每30秒发送一次心跳

    async def send_heartbeat(self):
        """发送心跳"""
        if not self.is_connected or not self.websocket:
            return
        try:
            heartbeat = {
                "type": "ping"
            }
            await self.websocket.send(json.dumps(heartbeat))
            logging.info(f"[{datetime.now()}] 发送心跳")
        except Exception as e:
            logging.error(f"发送心跳失败: {e}")
            self.is_connected = False

    async def handle_messages(self):
        """处理服务器消息"""
        while self.is_connected:
            try:
                message = await self.websocket.recv()
                data = json.loads(message)
                if data['type'] == 'task':
                    # 收到新任务
                    task = data['task']
                    logging.info(f"收到新任务: {task['title']}")
                    await self.process_task(task)
                elif data['type'] == 'pong':
                    # 心跳响应
                    logging.info(f"[{datetime.now()}] 收到心跳响应")
            except websockets.exceptions.ConnectionClosed:
                logging.warning("连接已关闭")
                self.is_connected = False
                break
            except Exception as e:
                logging.error(f"处理消息失败: {e}")
                break

    async def process_task(self, task):
        """处理任务"""
        logging.info(f"开始处理任务: {task['id']} - {task['desc']}")
        result = {"task_id": task["id"]}
        try:
            res = await task_mgr.build(
                ProcessTask(
                    id=str(task["id"]),
                    prompt=task["desc"],
                )
            )
            # 构造处理结果
            result.update({
                "error_msg": "",
                ** res
            })
            success = True
        except Exception as e:
            success = False
            result["error_msg"] = str(e)
        await self.send_result(task["id"], json.dumps(result), success)

    async def send_result(self, task_id, result, success):
        """发送任务结果到服务器"""
        if not self.is_connected or not self.websocket:
            return
        message = {
            "type": "result",
            "task_id": task_id,
            "result": result,
            "success": success,
        }
        try:
            # 发送前再次确认连接状态（避免并发关闭）
            if self.websocket:
                await self.websocket.send(json.dumps(message))
                logging.info(f"任务 {task_id} 处理完成，结果已发送")
            else:
                logging.warning(f"连接已关闭，无法发送任务 {task_id} 的结果")
        except websockets.exceptions.ConnectionClosed:
            logging.warning(f"连接已关闭，任务 {task_id} 结果发送失败")
        except Exception as e:
            logging.error(f"发送结果失败: {str(e)}")
            traceback.print_exc()

    async def run(self):
        """运行客户端"""
        if not await self.connect():
            return
        # 启动心跳和消息处理
        heartbeat_task = asyncio.create_task(self.start_heartbeat())
        message_task = asyncio.create_task(self.handle_messages())
        try:
            # 等待任务完成
            await asyncio.gather(heartbeat_task, message_task)
        except KeyboardInterrupt:
            logging.warning("客户端正在关闭...")
        finally:
            if self.websocket:
                await self.websocket.close()
